from .base import BaseDataset
import os

class Car(BaseDataset):
    nb_train_all = 8054
    nb_test_all = 8131
    def __init__(self, root, classes, transform=None):
        BaseDataset.__init__(self, root, classes, transform)

        # classes_train = range(0, 98)
        # classes_test = range(98, 196)

        if classes == range(0, 98):
            train = True
        else:
            train = False
        # if classes.start in classes_train:
        #     if classes.stop - 1 in classes_train:
        #         train = True
        #
        # if classes.start in classes_test:
        #     if classes.stop - 1 in classes_test:
        #         train = False

        self.I: list
        with open(
            os.path.join(
                root,
                '{}.txt'.format('train' if train else 'test')
            )
        ) as f:

            # f.readline()
            index = 0

            self.img2absid = {}
            for (path, class_id) in map(str.split, f):
                abspath = os.path.join(root, path)
                self.im_paths.append(abspath)
                self.img2absid[abspath] = index
                if train:
                    self.ys.append(int(class_id))
                else:
                    self.ys.append(int(class_id)+98)
                self.I.append(index)
                index += 1

            if train:
                assert index == type(self).nb_train_all
            else:
                assert index == type(self).nb_test_all
